#ifndef __WPE_PARSER__
#define __WPE_PARSER__

#include <string>
#include <vector>
#include <set>
#include <algorithm>
#include <stdio.h>

#include "thisinc.h"
#include "wpefio.h"
#include "wpefmt.h"

namespace wpe
{

	enum { WIN32_PAGE_SIZE = 0x1000 };

	static void fix_ms_dllname( char * p )
	{
		size_t len = strlen(p);
		if( !len ) return ;
		char * hp = p;
		while( isspace((unsigned char)*hp) ) ++ hp;
		char * tp = p + len - 1;
		while( tp > hp && isspace((unsigned char)*tp) ) -- tp;
		len = tp - hp + 1;
		if( len ) tcs::ncpy( p, hp, len );
		p[len] = 0;
		tcs::lwr( p );
	}
	static void fix_ms_dllname2( std::string & name, const char * p )
	{
		const char * pk = strrchr( p, '.' );
		if( !pk ) name = p;
		else name.assign( p, pk );
	}

	class PeParser : public RefalbeImp
	{
	public:

		typedef struct _impapi_t
		{
			uint32_t	iat;
			std::string	mod;
			std::string api;
		} impapi_t;

		typedef struct _expsym_t
		{
			uint32_t	ord;
			std::string	name;		// "" means exported by ord
		} expsym_t;

	protected:

		static bool less_impapi( const impapi_t & l, const impapi_t & r ) 
		{
			return l.iat < r.iat;
		}
		static bool less_name_ord_pair( const expsym_t & left, const expsym_t & right )
		{
			return left.name < right.name;
		}

		refp<FileObject> mfo;

		typedef std::vector<WPE_IMAGE_SECTION_HEADER>	TSectionTable;
		typedef std::vector<impapi_t>					TImportSet;
		typedef std::set<uint32_t>						TImpNameRvaSet;
		typedef std::vector<expsym_t>					TExportSet;
		typedef std::vector<uint32_t>					TExportRvaArray;

		uint32_t					mmach;
		TSectionTable				msections;
		TImportSet					miis32;
		TImpNameRvaSet				mnas32;
		TExportSet					mexpnams;
		TExportRvaArray				mexpeps;

	public:

		WPE_IMAGE_FILE_HEADER		mfh;
		WPE_IMAGE_OPTIONAL_HEADER32 moh32;
		WPE_IMAGE_OPTIONAL_HEADER64 moh64;

		size_t		is_x86_64()
		{
			return	mfh.Machine == WPE_IMAGE_FILE_MACHINE_AMD64 || 
					mfh.Machine == WPE_IMAGE_FILE_MACHINE_I386 ;
		}

		size_t		is_64bit()
		{
			return !( mfh.Characteristics & WPE_IMAGE_FILE_32BIT_MACHINE );
		}

		FileObject * get_fo()
		{
			return mfo;
		}
		size_t num_of_sections_with_header()
		{
			return msections.size();
		}
		WPE_IMAGE_SECTION_HEADER * section_header( size_t index )
		{
			if( index >= msections.size() ) return NULL;
			return &msections[index];
		}
		long open( FileObject * fo )
		{
			close();
			if( !fo ) return E_INVALIDARG;
			mfo = fo;
			return open__();
		}
		long close()
		{
			msections.clear();
			memset( &moh64, 0, sizeof(moh64) );
			memset( &moh32, 0, sizeof(moh32) );
			miis32.clear();
			mfo.dispose();
			return S_OK;
		}
		size_t num_of_imports()
		{
			return miis32.size();
		}
		impapi_t * get_impapi( size_t index )
		{
			if( index >= miis32.size() ) return NULL;
			return &miis32[index];
		}
		impapi_t * find_impapi( uint32_t iat )
		{
			impapi_t key; key.iat = iat;
			TImportSet::iterator it = std::lower_bound( miis32.begin(), miis32.end(), key, less_impapi );
			if( it == miis32.end() ) return NULL;
			if( it->iat != iat ) return NULL;
			return &(*it);
		}
		bool is_rva_of_import_name( uint32_t rva )
		{
			TImpNameRvaSet::iterator it = mnas32.find( rva );
			if( it == mnas32.end() ) return false;
			return true;
		}
		uint64_t get_image_base()
		{
			if( is_64bit() ) return moh64.ImageBase;
			return moh32.ImageBase;
		}
		uint32_t get_image_size()
		{
			if( is_64bit() ) return moh64.SizeOfImage;
			return moh32.SizeOfImage;
		}

		size_t get_num_of_exports()
		{
			return mexpeps.size();
		}

		long get_export_address( size_t index, uint32_t * ep )
		{
			if( index >= mexpeps.size() ) return E_FAIL;
			*ep = mexpeps[index];
			return S_OK;
		}

		ssize_t get_section_index_by_rva( uint32_t rva )
		{
			for( size_t i = 0; i < msections.size(); ++ i )
			{
				WPE_IMAGE_SECTION_HEADER & sh = msections[i];
				uint32_t sec_base = sh.VirtualAddress;
				uint32_t sec_size = sh.VirtualSize;
				if( rva < sec_base || rva >= sec_base + sec_size ) continue;
				return i;
			}
			return -1;
		}


	protected:

		long open__()
		{
			WPE_IMAGE_DOS_HEADER dh = {};
			long lr = readp_sc( 0, &dh, sizeof(dh) );
			if( lr < 0 ) return lr;

			if( dh.e_magic != WPE_IMAGE_DOS_SIGNATURE )
				return E_FAIL;

			uint32_t nfoff = (int32_t)dh.e_lfanew;

			le::U32 sig(0);
			lr = readp_sc( nfoff, &sig, sizeof(sig) );
			if( lr < 0 ) return lr;
			
			if( sig != WPE_IMAGE_NT_SIGNATURE ) return E_FAIL;

			lr = read_sc( &mfh, sizeof(mfh) );
			if( lr < 0 ) return lr;

			uint16_t ohs = mfh.SizeOfOptionalHeader;
			uint16_t sn = mfh.NumberOfSections;
			if( sn > 0x100 ) return E_FAIL;

			bool is_32 = (mfh.Characteristics & WPE_IMAGE_FILE_32BIT_MACHINE) != 0;

			if( is_32 )
			{
				lr = read_sc( &moh32, sizeof(moh32) );
				if( lr < 0 ) return lr;
				if( moh32.Magic != WPE_IMAGE_NT_OPTIONAL_HDR32_MAGIC ) return E_FAIL;
			}
			else
			{
				lr = read_sc( &moh64, sizeof(moh64) );
				if( lr < 0 ) return lr;
				if( moh64.Magic != WPE_IMAGE_NT_OPTIONAL_HDR64_MAGIC ) return E_FAIL;
			}

			size_t stoff = nfoff + sizeof(uint32_t) + sizeof(mfh) + mfh.SizeOfOptionalHeader;

			msections.resize( sn );
			size_t stbytes = sizeof(WPE_IMAGE_SECTION_HEADER)*sn;

			lr = readp_sc( stoff, msections.data(), stbytes );
			if( lr < 0 ) return lr;

			uint32_t prev_ve = 0;
			for( size_t i = 0; i < sn; ++ i )
			{
				WPE_IMAGE_SECTION_HEADER & hst = msections[i];
				hst.VirtualAddress = (uint32_t)__align_down((uint32_t)hst.VirtualAddress,WIN32_PAGE_SIZE);
				hst.VirtualSize = (uint32_t)__align_up((uint32_t)hst.VirtualSize,WIN32_PAGE_SIZE);
				if( prev_ve )
				{
					if( hst.VirtualAddress && hst.VirtualAddress < prev_ve )
						return E_FAIL;
				}
				prev_ve = hst.VirtualAddress + hst.VirtualSize;
			}

			size_t dde_num = is_64bit() ? moh64.NumberOfRvaAndSizes : moh32.NumberOfRvaAndSizes;

			WPE_IMAGE_SECTION_HEADER hst;
			memset( &hst, 0, sizeof(hst) );
			memcpy( hst.Name, "#header", 8 );
			hst.PointerToRawData = 0;
			hst.SizeOfRawData = stoff + sizeof(WPE_IMAGE_DATA_DIRECTORY)*dde_num + stbytes;
			hst.VirtualAddress = 0;
			hst.VirtualSize = sn ? msections[0].VirtualAddress : (uint32_t)__align_up((uint32_t)hst.SizeOfRawData,WIN32_PAGE_SIZE);

			msections.insert( msections.begin(), hst );

			RFAILED( parse_import__() );
			RFAILED( parse_export__() );

			return S_OK;
		}
		enum { MX_THUNKS_TO_PTOC = 5000 };
		enum { _MAX_PATH_ = 0x200 };

		template < class pe_thunk_t >
		long parse_import_desc__( WPE_IMAGE_IMPORT_DESCRIPTOR* piidesc, const char * libname )
		{
			typedef typename pe_thunk_t::thunk_value_t thunk_t;
			pe_thunk_t thunk;

			const thunk_t IMP_BY_ORD_BIT = (1LL<<(sizeof(thunk_t)*8-1));

			UTIL::sentry<char*> api_name( new char[_MAX_PATH_*2] );
			if( !api_name ) return E_OUTOFMEMORY;

			size_t eff = 0; long lr = 0;
			thunk_t thunk_rva = piidesc->OriginalFirstThunk;

			for( size_t i = 0; i < MX_THUNKS_TO_PTOC; ++ i, thunk_rva += sizeof(thunk) )
			{
				lr = imgread_sc( thunk_rva, &thunk, sizeof(thunk) );
				if( lr < 0 ) break;
				if( !thunk.MixData ) break;

				uint32_t iat = piidesc->FirstThunk + sizeof(thunk)*i;

				uint32_t rva = 0;
				if( thunk.MixData & IMP_BY_ORD_BIT )
				{
					std::string tmp;
					fix_ms_dllname2( tmp, libname );
					sprintf( api_name, "%s#%X", tmp.c_str(), thunk.MixData & (~IMP_BY_ORD_BIT) );
				}
				else
				{
					uint32_t api_rva = thunk.MixData + sizeof(uint16_t);
					lr = imgread( api_rva, api_name, _MAX_PATH_, &eff );
					if( lr < 0 || eff < 2 ) continue;
					api_name[eff-1] = 0;
					mnas32.insert( api_rva );
				}

				impapi_t tmp;
				tmp.iat = iat;
				tmp.mod = libname;
				tmp.api = api_name;

				miis32.push_back( tmp );
			}
			return S_OK;
		}
		
		enum { MAX_EXPORTS = 10000 };

		long parse_export__()
		{
			WPE_IMAGE_DATA_DIRECTORY dde;
			if( is_64bit() )
			{
				if( moh64.NumberOfRvaAndSizes <= WPE_IMAGE_DIRECTORY_ENTRY_EXPORT ) return S_OK;
				dde = moh64.DataDirectory[WPE_IMAGE_DIRECTORY_ENTRY_EXPORT];
			}
			else
			{
				if( moh32.NumberOfRvaAndSizes <= WPE_IMAGE_DIRECTORY_ENTRY_EXPORT ) return S_OK;
				dde = moh32.DataDirectory[WPE_IMAGE_DIRECTORY_ENTRY_EXPORT];
			}
			
			uint32_t ied_rva = dde.VirtualAddress;
			if( !ied_rva ) return S_OK;

			UTIL::sentry<char*> name_buffer( new char[_MAX_PATH_] );
			if( !name_buffer ) return E_OUTOFMEMORY;

			WPE_IMAGE_EXPORT_DIRECTORY ied;
			long lr = imgread_sc( ied_rva, &ied, sizeof(ied) );
			if( lr < 0 ) return E_FAIL;

			uint32_t rva_ord, rva_names, rva_eps;
			rva_ord = ied.AddressOfNameOrdinals;
			rva_names = ied.AddressOfNames;
			rva_eps = ied.AddressOfFunctions;

			uint32_t num_names = ied.NumberOfNames;
			uint32_t num_funcs = ied.NumberOfFunctions;

			if( num_names > MAX_EXPORTS || num_funcs > MAX_EXPORTS )
				return E_FAIL;

			mexpeps.resize(num_funcs);

			uint32_t export_ord_base = ied.Base;

			//////////////////////////////////////////////////////////////////////////
			size_t eff = 0;
			lr = imgread( rva_eps, mexpeps.data(), num_funcs * sizeof(uint32_t), &eff );
			if( lr < 0 ) return E_FAIL;

			num_funcs = eff/sizeof(uint32_t);
			mexpeps.resize( num_funcs );

			UTIL::sentry<char*> api_name_buf = new char[_MAX_PATH_];
			if( !api_name_buf ) return E_OUTOFMEMORY;

			size_t to_read = num_names * sizeof(uint32_t);
			for( size_t offset = 0; offset < to_read; offset += sizeof(uint32_t) )
			{
				uint16_t name_to_ord;
				lr = imgread( rva_ord + (offset>>1), &name_to_ord, sizeof(name_to_ord), &eff );
				if( lr < 0 ) return E_FAIL;

				if( eff != sizeof(name_to_ord) ) continue;
				if( name_to_ord >= num_funcs ) continue;

				uint32_t name_rva;

				lr = imgread( rva_names + offset, &name_rva, sizeof(uint32_t), &eff );

				if( lr < 0 ) return E_FAIL;
				if( eff != sizeof(uint32_t) ) continue;

				lr = imgread( name_rva, api_name_buf, _MAX_PATH_, &eff );
				if( lr < 0 || eff < 2 ) continue;

				api_name_buf[_MAX_PATH_-1] = 0;
				if( !strlen(api_name_buf) ) continue;

				expsym_t exp_pair;
				exp_pair.name = api_name_buf;
				exp_pair.ord = name_to_ord;
				mexpnams.push_back( exp_pair );
			}

			if( mexpnams.size() )
			{
				std::sort( mexpnams.begin(), mexpnams.end(), less_name_ord_pair );
			}

			return S_OK;

		}

		long parse_import__()
		{
			WPE_IMAGE_DATA_DIRECTORY dde;
			if( is_64bit() )
			{
				if( moh64.NumberOfRvaAndSizes <= WPE_IMAGE_DIRECTORY_ENTRY_IMPORT ) return S_OK;
				dde = moh64.DataDirectory[WPE_IMAGE_DIRECTORY_ENTRY_IMPORT];
			}
			else
			{
				if( moh32.NumberOfRvaAndSizes <= WPE_IMAGE_DIRECTORY_ENTRY_IMPORT ) return S_OK;
				dde = moh32.DataDirectory[WPE_IMAGE_DIRECTORY_ENTRY_IMPORT];
			}

			if( !dde.VirtualAddress ) 
				return S_OK;

			UTIL::sentry<char*> name_buffer( new char[_MAX_PATH_] );
			if( !name_buffer ) 
				return E_OUTOFMEMORY;

			WPE_IMAGE_IMPORT_DESCRIPTOR iid;
			size_t iid_rva = dde.VirtualAddress;

			for( size_t i = 0; i < 512; ++ i, iid_rva += sizeof(iid) )
			{
				long lr = imgread_sc( iid_rva, &iid, sizeof(iid) );
				if( lr < 0 ) 
					return E_FAIL;

				if( !iid.OriginalFirstThunk || !iid.FirstThunk )
					break;

				size_t eff = 0;
				lr = imgread( iid.Name, name_buffer, _MAX_PATH_, &eff );
				if( lr < 0 || eff < 2 ) 
					return E_FAIL;

				mnas32.insert( iid.Name );

				name_buffer[eff-1] = 0;

				fix_ms_dllname( name_buffer );

				if( is_64bit() )
					lr = parse_import_desc__<WPE_IMAGE_THUNK_DATA64>( &iid, name_buffer );
				else
					lr = parse_import_desc__<WPE_IMAGE_THUNK_DATA32>( &iid, name_buffer );
				if( lr < 0 ) 
					return lr;
			}

			std::sort( miis32.begin(), miis32.end(), less_impapi );

			return S_OK;
		}
		
	public:

		long rva_to_raw( uint32_t rva, uint32_t & off )
		{
			for( size_t i = 0; i < msections.size(); ++ i )
			{
				WPE_IMAGE_SECTION_HEADER & sh = msections[i];
				uint32_t sec_base = sh.VirtualAddress;
				uint32_t sec_size = sh.VirtualSize;
				if( rva < sec_base || rva >= sec_base + sec_size )
					continue;
				if( !sh.SizeOfRawData ) return E_FAIL;
				off = rva - sh.VirtualAddress + sh.PointerToRawData;
				return S_OK;
			}
			return E_FAIL;
		}
		long imgread( uint32_t rva, void * buffer, size_t bytes, size_t * eff )
		{
			uint32_t off = 0;
			long lr = rva_to_raw( rva, off );
			if( lr < 0 ) return lr;
			return readp_( off, buffer, bytes, eff );
		}
		long imgread_sc( uint32_t rva, void * buffer, size_t bytes )
		{
			size_t eff = 0;
			long lr = imgread( rva, buffer, bytes, &eff );
			if( lr < 0 ) return lr;
			if( eff != bytes ) return E_FAIL;
			return S_OK;
		}
		long read_( void * buffer, size_t bytes, size_t * eff )
		{
			return mfo->read( buffer, bytes, eff );
		}
		long read_sc( void * buffer, size_t bytes )
		{
			size_t eff = 0;
			long lr = mfo->read( buffer, bytes, &eff );
			if( lr < 0 ) return lr;
			if( eff != bytes ) return E_FAIL;
			return S_OK;
		}
		long readp_( uint64_t pos, void * buffer, size_t bytes, size_t * eff )
		{
			return mfo->readp( pos, buffer, bytes, eff );
		}
		long readp_sc( uint64_t pos, void * buffer, size_t bytes )
		{
			size_t eff = 0;
			long lr = readp_( pos, buffer, bytes, &eff );
			if( lr < 0 ) return lr;
			if( eff != bytes ) return E_FAIL;
			return S_OK;
		}
	};








};



#endif